// Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Basic unittest for general_installer

#include "memento_common.h"
#include "general_installer.h"

#include <gtest/gtest.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <stdint.h>
#include <iostream>

using std::string;

namespace chromeos_memento_updater {

const string kProgramDir = ".";

class GeneralInstallerTest : public ::testing::Test {
 protected:
  GeneralInstallerTest()
      : installer_(kProgramDir, "") {
  }
  GeneralInstaller installer_;
};

TEST_F(GeneralInstallerTest, InstallSingleImage) {
  TempfileHandler input;
  TempfileHandler output;
  char test_string[] = "This is a test file\n";
  string input_path = input.GetFilePath();
  string output_path = output.GetFilePath();
  FILE* fp = fopen(input_path.c_str(), "w");
  fprintf(fp, "%s", test_string);
  fclose(fp);
  fp = fopen(input_path.c_str(), "r");
  EXPECT_TRUE(installer_.InstallImage(output_path, string(), fp) == 0);
  fclose(fp);

  fp = fopen(output_path.c_str(), "r");
  char buffer[100];
  fgets(buffer, 100 ,fp);
  EXPECT_TRUE(strcmp(buffer, test_string) == 0);
  EXPECT_TRUE(fgetc(fp) == EOF);
  fclose(fp);
}

TEST_F(GeneralInstallerTest, InstallTwoImages) {
  TempfileHandler input;
  TempfileHandler output[2];
  char test_string[2][40] = {"This is the first test\n",
                             "This is the second test\n"};
  string input_path = input.GetFilePath();
  string output_path[2];
  for (int i = 0; i < 2; i++)
    output_path[i] = output[i].GetFilePath();

  // Prepare input file
  FILE* fp = fopen(input_path.c_str(), "w");
  int64_t first_file_size = strlen(test_string[0]);
  first_file_size = htobe64(first_file_size);
  fwrite(&first_file_size, sizeof(first_file_size), 1, fp);
  for (int i = 0; i < 2; i++)
    fprintf(fp, "%s", test_string[i]);
  fclose(fp);

  // Use input file to write device
  fp = fopen(input_path.c_str(), "r");
  EXPECT_TRUE(installer_.InstallImage(output_path[1], output_path[0], fp) ==0);
  fclose(fp);

  // Check if output is correct
  for (int i = 0; i < 2; i++) {
    fp = fopen(output_path[i].c_str(), "r");
    char buffer[100];
    fgets(buffer, 100 ,fp);
    EXPECT_TRUE(strcmp(buffer, test_string[i]) == 0);
    EXPECT_TRUE(fgetc(fp) == EOF);
    fclose(fp);
  }
}

TEST_F(GeneralInstallerTest, ZeroOutDevice) {
  TempfileHandler handler;
  string file_path = handler.GetFilePath();
  FILE* fp = fopen(file_path.c_str(), "w");

  // print something so that the file is not empty
  for (int i = 0; i < 300 ; i++) {
    fprintf(fp, "test");
  }
  fclose(fp);
  installer_.ZeroOutDevice(file_path);
  fp = fopen(file_path.c_str(), "r");
  char buffer[4096];
  fgets(buffer, 4096, fp);

  // Buffer should be empty
  for (int i = 0; i<4096; i++)
    EXPECT_FALSE(buffer[i]);
  fclose(fp);
}

TEST_F(GeneralInstallerTest, UpdateFirmware) {
  TempfileHandler handler;
  string file_path = handler.GetFilePath();
  string loop_device = "/dev/loop0";  // TODO(hungte) don't hard-code loop0.

  // Create an fake image
  ASSERT_TRUE(system(("dd if=/dev/zero of=" + file_path +
                      " bs=1024 count=5000 >/dev/null").c_str()) == 0);
  ASSERT_TRUE(system(("/sbin/mkfs.ext2 -F " + file_path +
                      " >/dev/null").c_str()) == 0);
  ASSERT_TRUE(system(("losetup " + loop_device +
                      " " + file_path).c_str()) == 0);
  {
    MountHandler mount_handler(loop_device, "ext2", 0);
    string mount_point = mount_handler.GetMountPoint();

    // Recursively make directory
    ASSERT_TRUE(mkdir((mount_point + "/usr").c_str(), S_IRWXU | S_IRWXO) == 0);
    ASSERT_TRUE(mkdir((mount_point + "/usr/sbin").c_str(), S_IRWXU | S_IRWXO)
                == 0);

    FILE* fp = fopen((mount_point +
                      "/usr/sbin/chromeos-firmwareupdate").c_str(), "w");
    ASSERT_TRUE(fp != NULL);
    fprintf(fp, "echo \"Firmware Updated.\"");
    fclose(fp);
  }

  // Program will exit when fail
  string updater_path = installer_.PrepareFirmware(loop_device);

  // Program will exit if fail to run script
  installer_.UpdateFirmware(updater_path, "");

  ASSERT_TRUE(system(("losetup -d " + loop_device).c_str()) == 0);
}

TEST_F(GeneralInstallerTest, ActivateImage) {
  TempfileHandler file_handler;
  string file_path = file_handler.GetFilePath();
  string loop_device = "/dev/loop0";
  // Create an fake image
  ASSERT_TRUE(system(("dd if=/dev/zero of=" + file_path +
                      " bs=1024 count=5000 >/dev/null").c_str()) == 0);
  ASSERT_TRUE(system(("/sbin/mkfs.ext2 -F " + file_path +
                      " >/dev/null").c_str()) == 0);
  ASSERT_TRUE(system(("losetup " + loop_device +
                      " " + file_path).c_str()) == 0);

  // Create an fake postinst
  MountHandler* mount_handler = new MountHandler(loop_device, "ext2", 0);
  string mount_point = mount_handler->GetMountPoint();
  FILE* fp = fopen((mount_point + "/postinst").c_str(), "w");
  fprintf(fp, "echo \"Image Activated.\"");
  fclose(fp);
  chmod((mount_point + "/postinst").c_str(),
        S_IXUSR | S_IXGRP | S_IXOTH);
  // Make sure destructor is called before ActivateImage()
  delete mount_handler;
  // Program will exit if fail to run script
  installer_.ActivateImage(loop_device);

  ASSERT_TRUE(system(("losetup -d " + loop_device).c_str()) == 0);
}

}  // namespace chromeos_memento_updater

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}
